Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

deepmd/dpmodel/utils/env_mat_stat.py (1)

77-78: Potential KeyError if link_stats has mismatched keys.

The iteration assumes base_stats and link_stats have identical keys. If link_obj.stats has a different set of keys (e.g., different type counts or missing entries), accessing link_stats[kk] will raise a KeyError.

Consider adding defensive handling:

🛡️ Proposed defensive fix

     for kk in base_stats:
+        if kk not in link_stats:
+            raise ValueError(f"Stats key '{kk}' missing in link_obj, cannot merge.")
         merged_stats[kk] = base_stats[kk] + link_stats[kk] * model_prob

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/utils/env_mat_stat.py` around lines 77 - 78, The loop in
env_mat_stat.py assumes link_stats contains every key in base_stats, which can
raise KeyError; update the loop in the function that builds merged_stats so it
defensively fetches link_stats values (use link_stats.get(kk, default)) and
choose a sensible default that matches base_stats' shape/type (e.g., 0 for
scalars or numpy.zeros_like(base_stats[kk]) for arrays) before computing
merged_stats[kk] = base_stats[kk] + link_val * model_prob; reference the
variables merged_stats, base_stats, link_stats, model_prob to locate and change
the code.

deepmd/pt_expt/descriptor/hybrid.py (1)

33-39: Rename unused loop variable des to _des.

The loop variable is not used within the loop body. Per Ruff B007, prefix it with an underscore.

🔧 Suggested fix

-        for ii, des in enumerate(self.descrpt_list):
+        for ii, _des in enumerate(self.descrpt_list):
             self.descrpt_list[ii].share_params(

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/pt_expt/descriptor/hybrid.py` around lines 33 - 39, The loop in
hybrid.py uses an unused loop variable named des; update the for-loop header to
use _des instead (for ii, _des in enumerate(self.descrpt_list):) so it complies
with Ruff B007, leaving the body that calls
self.descrpt_list[ii].share_params(base_class.descrpt_list[ii], shared_level,
model_prob=model_prob, resume=resume) unchanged; ensure only the loop variable
name is changed and no other logic is modified.

source/tests/pt_expt/descriptor/test_dpa3.py (1)

268-270: Consider using underscore prefix for unused unpacked variables.

nf and nloc from self.nlist.shape are not used in this test. Per Ruff RUF059, prefix them with underscores.

🔧 Suggested fix

-        nf, nloc, nnei = self.nlist.shape
+        _nf, _nloc, nnei = self.nlist.shape

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@source/tests/pt_expt/descriptor/test_dpa3.py` around lines 268 - 270, The
unpacking of self.nlist.shape assigns nf and nloc which are unused; change the
unpack to use underscore-prefixed names (e.g., _nf, _nloc, nnei) so only nnei
remains a meaningful variable while satisfying RUF059; update the tuple
assignment where self.nlist.shape is unpacked in the test (currently "nf, nloc,
nnei = self.nlist.shape") to use "_nf, _nloc, nnei = self.nlist.shape".

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

deepmd/pt_expt/entrypoints/main.py (2)

292-306: ⚠️ Potential issue | 🟠 Major

Always destroy the process group in a finally.

If get_trainer() or trainer.run() raises after Line 293, destroy_process_group() is skipped and the worker stays in a bad distributed state.

Minimal fix

-    if os.environ.get("LOCAL_RANK") is not None:
-        dist.init_process_group(backend="cuda:nccl,cpu:gloo")
-
-    trainer = get_trainer(
-        config,
-        init_model,
-        restart,
-        finetune_model=finetune,
-        finetune_links=finetune_links,
-        shared_links=shared_links,
-    )
-    trainer.run()
-
-    if dist.is_available() and dist.is_initialized():
-        dist.destroy_process_group()
+    try:
+        if os.environ.get("LOCAL_RANK") is not None:
+            dist.init_process_group(backend="cuda:nccl,cpu:gloo")
+
+        trainer = get_trainer(
+            config,
+            init_model,
+            restart,
+            finetune_model=finetune,
+            finetune_links=finetune_links,
+            shared_links=shared_links,
+        )
+        trainer.run()
+    finally:
+        if dist.is_available() and dist.is_initialized():
+            dist.destroy_process_group()

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/pt_expt/entrypoints/main.py` around lines 292 - 306, Wrap the
distributed init/destroy in a try/finally so the process group is always
cleaned: call dist.init_process_group(...) (existing code) before creating the
trainer, then surround get_trainer(...) and trainer.run() with try/finally and
move the dist.destroy_process_group() into the finally block; keep the same
guards using dist.is_available() and dist.is_initialized() before calling
dist.destroy_process_group() to avoid errors if init failed.

---

88-95: ⚠️ Potential issue | 🟠 Major

Create nested stat directories in single-task mode too.

This branch still uses plain mkdir(), so a nested path like ./stat_files/model_1 fails in single-task mode even though the multi-task branch already handles parent creation.

Minimal fix

-                    Path(stat_file_path).mkdir()
+                    Path(stat_file_path).mkdir(parents=True, exist_ok=True)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/pt_expt/entrypoints/main.py` around lines 88 - 95, The current branch
that creates stat_file_path in single-task mode uses
Path(stat_file_path).mkdir() without parents=True, so nested directories (e.g.,
"./stat_files/model_1") fail; update the creation logic where stat_file_path is
checked (the block that tests Path(stat_file_path).exists() and then calls
Path(stat_file_path).mkdir()) to create parent directories as needed by calling
mkdir with parents=True and exist_ok=True (or equivalent), ensuring the
DPPath(stat_file_path, "a") line still executes afterwards; keep the HDF5 file
branch unchanged.

source/tests/pt_expt/test_training_ddp.py (2)

62-66: ⚠️ Potential issue | 🟠 Major

Replace free-port probing with a persistent rendezvous.

This helper closes the socket before the workers call init_process_group(), so another process can grab the port in between and make the DDP suite flaky under parallel CI. A file:// rendezvous or a pre-created TCPStore is safer here.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@source/tests/pt_expt/test_training_ddp.py` around lines 62 - 66, The
_find_free_port helper closes the socket before workers call init_process_group,
allowing a race where another process can take the port; replace this free-port
probing with a persistent rendezvous mechanism (e.g., use a file:// rendezvous
URL or create and reuse a pre-created torch.distributed.TCPStore) so the
port/address remains reserved across worker startup. Update the test setup that
calls _find_free_port to instead create a stable rendezvous target (or a shared
TCPStore) and pass that rendezvous information into init_process_group (or pass
the initialized TCPStore into the processes) to avoid the transient socket race.

---

622-631: ⚠️ Potential issue | 🟠 Major

Add a hard 60s timeout around each spawned training scenario.

A stuck rendezvous or worker here will hang the suite indefinitely because every test calls mp.spawn(..., join=True) with no outer timeout. Please wrap each spawned run in a watchdog process (or equivalent timeout mechanism) so the test fails fast instead of wedging CI. As per coding guidelines, **/tests/**/*training*.py: Set training test timeouts to 60 seconds maximum for validation purposes, as real training takes hours or days.

Also applies to: 662-671, 708-713, 815-820, 879-892, 938-951, 1053-1058, 1097-1102, 1195-1200, 1549-1561

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@source/tests/pt_expt/test_training_ddp.py` around lines 622 - 631, Wrap the
mp.spawn(...) call inside a watchdog multiprocessing.Process to enforce a 60s
hard timeout: create a Process whose target executes
mp.spawn(_worker_single_task_train, args=(2, port, self.data_dir, result_dict),
nprocs=2, join=True), start it, call p.join(60), and if p.is_alive() then
p.terminate(), p.join(), and raise an AssertionError (or fail the test)
indicating a timeout; apply the same pattern for other tests that call mp.spawn
so any stuck rendezvous/workers are killed and the test fails fast. Ensure you
reference the same symbols (_worker_single_task_train and mp.spawn) and perform
cleanup (terminate + join) before failing.

source/tests/pt_expt/fitting/test_fitting_stat.py (1)

148-150: ⚠️ Potential issue | 🟡 Minor

Extend the data preflight to cover data_1 as well.

test_sharefitting_using_default_fparam() later reads _PT_DATA_NO_FPARAM, but the shared skip helper never checks that directory. On workers without pt/water/data/data_1, this turns into a FileNotFoundError instead of a clean skip.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@source/tests/pt_expt/fitting/test_fitting_stat.py` around lines 148 - 150,
The preflight helper _skip_if_no_data currently only checks _PT_DATA; update it
to also check the secondary dataset used in tests (_PT_DATA_NO_FPARAM or the
data_1 directory) and raise unittest.SkipTest if either required directory is
missing. In practice, modify _skip_if_no_data to verify os.path.isdir(_PT_DATA)
and os.path.isdir(_PT_DATA_NO_FPARAM) (or the path for data_1) and produce a
clear SkipTest message naming the missing directory so
test_sharefitting_using_default_fparam() will be skipped instead of raising
FileNotFoundError.

deepmd/pt_expt/train/training.py (1)

1226-1239: ⚠️ Potential issue | 🟠 Major

Don't advance training data just to log other tasks.

This rank-0-only logging path pulls extra training batches for the non-selected heads, so disp_freq changes the actual training stream and distributed runs drift further because only rank 0 consumes those iterators. Please switch these display-only metrics to validation data or a cached/non-advancing sample.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/pt_expt/train/training.py` around lines 1226 - 1239, The loop under
self.model_keys currently pulls training batches for non-selected heads
(self.get_data(is_train=True, task_key=_key)) which advances the training
iterator on rank 0; change it to use validation/non-advancing samples instead
(e.g., call self.get_data(is_train=False, task_key=_key) or fetch from a cached
display sample like self._display_samples[_key] if available) and avoid calling
optimizer.zero_grad on these display-only queries; keep using _unwrapped(...) to
compute metrics but ensure the data source does not advance training iterators
so train_results is populated without consuming training batches.

deepmd/pt_expt/fitting/invar_fitting.py (1)

45-77: ⚠️ Potential issue | 🟠 Major

Seed the base fitting stats before aliasing shared buffers.

If the first shared head has empty fparam/aparam stats and a later head is the first one with real data, these branches skip the merge and still alias base_class's buffers into every branch. The resulting shared stats become model-order dependent and can drop the first populated stats entirely. Please initialize base_class._param_stats[...] and the corresponding avg/inv-std buffers from self_stats when base_stats is empty, then continue chaining merges from that seeded state.

Also applies to: 83-115

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/pt_expt/fitting/invar_fitting.py` around lines 45 - 77, When
base_class has empty stats for "fparam" (base_stats == []), seed
base_class._param_stats["fparam"] with a deep copy of self_stats and initialize
base_class.fparam_avg and base_class.fparam_inv_std buffers from self_stats'
computed averages/stds (using compute_avg/compute_std with protection and
matching dtypes/device) before performing the merge logic that uses numb_fparam,
model_prob and writes back merged stats; this ensures subsequent chaining uses
the seeded state. Update the same pattern for the "aparam" branch as well (seed
base_class._param_stats["aparam"] and base_class.aparam_avg/aparam_inv_std from
self.get_param_stats()["aparam"] when base_stats empty) so aliasing via
self._buffers[...] always reflects the first non-empty head.

source/tests/pt_expt/test_multitask.py (2)

1951-1952: ⚠️ Potential issue | 🟡 Minor

Unused loop variable n2 should be marked intentional.

Rename the unused variable to _ (or _n2) to satisfy Ruff.

Suggested fix

-        for (n1, p1), (n2, p2) in zip(
+        for (n1, p1), (_, p2) in zip(
             mt_desc.named_parameters(), mt_desc_2.named_parameters(), strict=True
         ):

As per coding guidelines, "Install linter and run ruff check . before committing changes or the CI will fail".

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@source/tests/pt_expt/test_multitask.py` around lines 1951 - 1952, The loop
over mt_desc.named_parameters() and mt_desc_2.named_parameters() declares an
unused variable n2 which triggers the linter; rename n2 to _ (or _n2) in the
for-loop header that iterates over (n1, p1), (n2, p2) returned from
mt_desc.named_parameters() and mt_desc_2.named_parameters() so the unused name
is explicit and Ruff will be satisfied while keeping p2 used as before.

---

945-947: ⚠️ Potential issue | 🟡 Minor

Unused finetune_links_true still triggers Ruff and will fail CI.

This binding is unused in the test and should be explicitly discarded.

Suggested fix

-        model_config_true, finetune_links_true = get_finetune_rules(
+        model_config_true, _ = get_finetune_rules(
             ckpt_path, deepcopy(ft_config_true["model"]), change_model_params=True
         )

As per coding guidelines, "Install linter and run ruff check . before committing changes or the CI will fail".

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@source/tests/pt_expt/test_multitask.py` around lines 945 - 947, The test
binds an unused value finetune_links_true from get_finetune_rules, which
triggers the linter; update the call site in test_multitask.py so the unused
return is explicitly discarded (e.g., assign to _ or use Python's unpacking like
model_config_true, _ = get_finetune_rules(...)) while preserving the existing
model_config_true and change_model_params=True arguments; ensure you only change
the left-hand side of the assignment where get_finetune_rules(...) is called.

deepmd/pt_expt/descriptor/hybrid.py (1)

18-41: LGTM with minor nit: unused loop variable.

The share_params implementation correctly delegates to each sub-descriptor. However, the loop variable des on line 33 is extracted but never used—only the index ii is needed.

🔧 Proposed fix to rename unused variable

         if shared_level == 0:
-            for ii, des in enumerate(self.descrpt_list):
+            for ii, _des in enumerate(self.descrpt_list):
                 self.descrpt_list[ii].share_params(

Alternatively, iterate directly with range(len(self.descrpt_list)):

         if shared_level == 0:
-            for ii, des in enumerate(self.descrpt_list):
+            for ii in range(len(self.descrpt_list)):
                 self.descrpt_list[ii].share_params(

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/pt_expt/descriptor/hybrid.py` around lines 18 - 41, The for-loop in
share_params uses an unused loop variable `des` (in the method share_params of
the hybrid descriptor class), so replace the enumerate form with either a
range-based loop (for ii in range(len(self.descrpt_list))) or use underscore for
the unused value (for ii, _ in enumerate(self.descrpt_list)) to avoid the
unused-variable warning; keep the body calling
self.descrpt_list[ii].share_params(...) with the same arguments.

deepmd/dpmodel/utils/env_mat_stat.py (1)

88-109: Consider handling case where neither davg/dstd nor mean/stddev attributes exist.

The function silently does nothing if neither attribute pair exists, which could lead to subtle bugs where stats are computed but buffers aren't updated. A warning or explicit handling may help surface configuration issues.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/utils/env_mat_stat.py` around lines 88 - 109, The code
currently updates buffers on base_obj when it has davg/dstd or mean/stddev but
silently skips if neither pair exists; add explicit handling for that case in
the same function: after the existing if/elif on hasattr(base_obj, "davg") /
hasattr(base_obj, "mean"), add an else branch that logs a warning (or raises an
exception depending on desired strictness) referencing base_obj and the missing
attributes, e.g., mentioning that neither davg/dstd nor mean/stddev were found
and stats were not applied; ensure the logger or exception uses identifiers like
base_obj, davg, dstd, mean, stddev, and respects set_davg_zero semantics if
relevant.

source/tests/pt_expt/descriptor/test_dpa3.py (1)

264-313: LGTM with minor nit: unused unpacked variables.

The test correctly validates sharing semantics for DPA3. The unpacked variables nf and nloc on line 268 are unused—only nnei is needed.

🔧 Proposed fix

-        nf, nloc, nnei = self.nlist.shape
+        _, _, nnei = self.nlist.shape

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@source/tests/pt_expt/descriptor/test_dpa3.py` around lines 264 - 313, In
test_share_params, remove the unused unpacked variables nf and nloc from the
assignment of self.nlist.shape; only extract nnei (e.g., use a single-element
assignment or ignore the first two values with placeholders) so the test no
longer defines unused variables; update the line in test_share_params that
currently reads "nf, nloc, nnei = self.nlist.shape" to only capture nnei
(referencing the test_share_params function and the self.nlist.shape tuple).

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

deepmd/dpmodel/fitting/general_fitting.py (1)

413-415: Consider returning a defensive copy from get_param_stats

GeneralFitting.get_param_stats() currently returns internal mutable state directly. Returning a shallow copy (or read-only mapping) would reduce accidental external mutation of GeneralFitting._param_stats.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/fitting/general_fitting.py` around lines 413 - 415,
get_param_stats currently returns the internal mutable attribute _param_stats
directly which risks external mutation; change GeneralFitting.get_param_stats to
return a defensive copy or read-only view (e.g., return dict(self._param_stats)
for a shallow copy or wrap with types.MappingProxyType for an immutable mapping)
so callers cannot accidentally modify the internal _param_stats; ensure the
method still returns the same typing dict[str, list[StatItem]] and keep behavior
when _param_stats is missing by copying the fallback {}.

deepmd/pt_expt/train/training.py (3)

901-903: ⚠️ Potential issue | 🔴 Critical

Use _unwrapped.model inside _compile_model().

After DDP wrapping, self.wrapper is a DistributedDataParallel, so self.wrapper.model[task_key] here will fail when compile is enabled in distributed mode. Read and write through self._unwrapped.model[...] instead.

Minimal fix

         for task_key in self.model_keys:
-            model = self.wrapper.model[task_key]
+            model = self._unwrapped.model[task_key]
...
-            self.wrapper.model[task_key] = _CompiledModel(
+            self._unwrapped.model[task_key] = _CompiledModel(
                 model, compiled_lower, max_nall, task_compile_opts
             )

#!/bin/bash
# Verify that DDP wrapping happens before _compile_model(), and that
# _compile_model() still dereferences self.wrapper.model instead of self._unwrapped.model.

sed -n '645,668p' deepmd/pt_expt/train/training.py
printf '\n---\n'
sed -n '901,1005p' deepmd/pt_expt/train/training.py
printf '\n---\n'
sed -n '1074,1079p' deepmd/pt_expt/train/training.py

Expected result: the first block shows self.wrapper being replaced by DistributedDataParallel, the second block shows direct self.wrapper.model[...] access inside _compile_model(), and the third block shows the existing _unwrapped helper that should be used instead.

Also applies to: 1003-1005

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/pt_expt/train/training.py` around lines 901 - 903, Change
_compile_model to access model modules via self._unwrapped.model instead of
self.wrapper.model: iterate over self.model_keys and use
self._unwrapped.model[task_key] when reading or assigning models (replace
occurrences of self.wrapper.model[...] in _compile_model). This ensures
compatibility after DistributedDataParallel wrapping; search for _compile_model,
self.wrapper, and self._unwrapped.model to update all direct wrapper.model
dereferences (also apply same change around lines referenced for the second
occurrence).

---

638-642: ⚠️ Potential issue | 🟡 Minor

Make both probability-map zips strict.

Both dict(zip(...)) calls still silently truncate if self.model_keys and self.model_prob ever drift, and Ruff already flags both sites.

Minimal fix

-                model_key_prob_map=dict(zip(self.model_keys, self.model_prob)),
+                model_key_prob_map=dict(
+                    zip(self.model_keys, self.model_prob, strict=True)
+                ),

-                    model_key_prob_map=dict(zip(self.model_keys, self.model_prob)),
+                    model_key_prob_map=dict(
+                        zip(self.model_keys, self.model_prob, strict=True)
+                    ),

As per coding guidelines, **/*.py: Install linter and run ruff check . before committing changes or the CI will fail.

Also applies to: 846-849

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/pt_expt/train/training.py` around lines 638 - 642, Summary: Replace
silent-truncating dict(zip(...)) usages with strict length checks to prevent
silent loss when self.model_keys and self.model_prob drift. Fix: before calling
self.wrapper.share_params(...) where
model_key_prob_map=dict(zip(self.model_keys, self.model_prob)) (and the other
similar call later), validate that len(self.model_keys) == len(self.model_prob)
and raise a clear ValueError if not; then construct model_key_prob_map using
dict(zip(...)) safely. Reference symbols: self.wrapper.share_params,
model_key_prob_map, self.model_keys, self.model_prob; apply the same validation
and construction at the other occurrence flagged in the review.

---

1225-1235: ⚠️ Potential issue | 🟠 Major

Per-task logging still advances the training iterators.

On rank 0, this logging path fetches get_data(is_train=True, task_key=_key) for non-active heads just to populate metrics. That makes disp_freq change which batches those heads actually train on, and only rank 0 consumes those extra batches.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/pt_expt/train/training.py` around lines 1225 - 1235, The per-task
logging loop is advancing training iterators because it calls
get_data(is_train=True, task_key=_key) for non-active heads; change that so
non-active heads do not consume training batches—either skip the
get_data/_unwrapped call for _key != task_key or call get_data(is_train=False,
task_key=_key) (or another non-consuming/validation accessor) so the training
iterator is not advanced; update the loop around self.model_keys (the block
using self.optimizer.zero_grad(), self.get_data(...), and self._unwrapped(...))
to only fetch training data for the active task_key and use a non-consuming path
for logging metrics.

source/tests/pt_expt/test_multitask.py (1)

939-941: ⚠️ Potential issue | 🟡 Minor

Drop the unused test bindings before Ruff does.

finetune_links_true and n2 are never read, so this file still trips the required lint step.

Minimal fix

-        model_config_true, finetune_links_true = get_finetune_rules(
+        model_config_true, _ = get_finetune_rules(
             ckpt_path, deepcopy(ft_config_true["model"]), change_model_params=True
         )

-        for (n1, p1), (n2, p2) in zip(
+        for (n1, p1), (_, p2) in zip(
             mt_desc.named_parameters(), mt_desc_2.named_parameters(), strict=True
         ):

As per coding guidelines, **/*.py: Install linter and run ruff check . before committing changes or the CI will fail.

Also applies to: 1945-1947

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@source/tests/pt_expt/test_multitask.py` around lines 939 - 941, The test
creates unused bindings (finetune_links_true and n2) from get_finetune_rules and
another unpack that Ruff flags; remove those unused variables by changing the
unpack to only capture what is used (e.g., assign just model_config_true =
get_finetune_rules(...)[0] or unpack as model_config_true, _ =
get_finetune_rules(...)) and similarly drop or replace n2 with an underscore in
the other unpack; update the lines that call get_finetune_rules and any other
tuple unpacking to avoid creating unused names.

source/tests/pt_expt/test_training_ddp.py (3)

622-631: ⚠️ Potential issue | 🟠 Major

Bound each spawned training case to 60 seconds.

All of these mp.spawn(..., join=True) paths can hang indefinitely on rendezvous or worker failures. Please route them through a helper that enforces a 60s timeout and fails the test if the subprocess stays alive.

As per coding guidelines, **/tests/**/*training*.py: Set training test timeouts to 60 seconds maximum for validation purposes, as real training takes hours or days.

Also applies to: 666-671, 708-713, 815-820, 879-892, 938-951, 1053-1058, 1097-1102, 1195-1200, 1549-1560

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@source/tests/pt_expt/test_training_ddp.py` around lines 622 - 631, Wrap the
mp.spawn call in test_ddp_single_task_trains with a helper that enforces a 60s
timeout and fails the test if any spawned subprocess remains alive; specifically
replace the direct mp.spawn(...) invocation that launches
_worker_single_task_train with a timeout-aware runner (e.g.,
run_spawn_with_timeout(spawn_fn=_worker_single_task_train, nprocs=2, args=(2,
port, self.data_dir, result_dict), timeout=60)) so rendezvous or worker hangs
are terminated and the test asserts failure. Apply the same pattern to the other
mp.spawn sites noted (the ranges referenced in the comment) so every training
test uses the common 60s timeout helper.

---

33-47: ⚠️ Potential issue | 🔴 Critical

DEVICE=cpu is set too late for spawned workers.

With mp.spawn, each child imports this module before executing _worker_*, so these module-level deepmd.pt_expt imports can resolve deepmd.pt_expt.utils.env.DEVICE before the worker sets the environment. That makes the test behavior depend on the parent process environment instead of staying CPU-only.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@source/tests/pt_expt/test_training_ddp.py` around lines 33 - 47, The test
sets DEVICE too late so module-level imports (e.g., importing get_trainer,
get_finetune_rules, preprocess_shared_params, normalize, update_deepmd_input)
cause deepmd.pt_expt to read env.DEVICE before spawned workers set it; to fix,
ensure DEVICE is set to "cpu" before those imports by either moving
os.environ["DEVICE"]="cpu" to the top of the module (before any deepmd.pt_expt
imports) or by deferring these imports into the worker entrypoints (the worker
functions named like _worker_*) so the environment is configured first.

---

62-66: ⚠️ Potential issue | 🟠 Major

The free-port probe is still racy.

_find_free_port() closes the socket before the workers call init_process_group(), so another process can claim the port in between and make this suite flaky under parallel CI.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@source/tests/pt_expt/test_training_ddp.py` around lines 62 - 66, The current
_find_free_port() closes the probe socket before workers call
init_process_group(), which allows another process to grab the port; change
_find_free_port to bind and listen but not close the socket (e.g., set
SO_REUSEADDR, bind(("127.0.0.1", 0)), listen()) and return the open socket and
its port number so the test can hold the socket open until after
init_process_group() completes; update callers to keep the returned socket alive
(and only close it after init_process_group() on all workers) to eliminate the
race.

deepmd/pt_expt/entrypoints/main.py (2)

290-304: ⚠️ Potential issue | 🟠 Major

Always destroy the process group in a finally.

If get_trainer() or trainer.run() raises, teardown is skipped and the process group stays initialized in a bad state.

Minimal fix

-    if os.environ.get("LOCAL_RANK") is not None:
-        dist.init_process_group(backend="cuda:nccl,cpu:gloo")
-
-    trainer = get_trainer(
-        config,
-        init_model,
-        restart,
-        finetune_model=finetune,
-        finetune_links=finetune_links,
-        shared_links=shared_links,
-    )
-    trainer.run()
-
-    if dist.is_available() and dist.is_initialized():
-        dist.destroy_process_group()
+    try:
+        if os.environ.get("LOCAL_RANK") is not None:
+            dist.init_process_group(backend="cuda:nccl,cpu:gloo")
+
+        trainer = get_trainer(
+            config,
+            init_model,
+            restart,
+            finetune_model=finetune,
+            finetune_links=finetune_links,
+            shared_links=shared_links,
+        )
+        trainer.run()
+    finally:
+        if dist.is_available() and dist.is_initialized():
+            dist.destroy_process_group()

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/pt_expt/entrypoints/main.py` around lines 290 - 304, Wrap the code
that calls dist.init_process_group, get_trainer, and trainer.run in a
try/finally so the process group is always torn down; specifically, after
calling dist.init_process_group (when os.environ.get("LOCAL_RANK") is not None)
call get_trainer(...) and trainer.run() inside a try block and in the finally
call dist.destroy_process_group() only if dist.is_available() and
dist.is_initialized(); update the block around dist.init_process_group,
get_trainer, and trainer.run to ensure cleanup even if get_trainer() or
trainer.run() raises.

---

87-95: ⚠️ Potential issue | 🟠 Major

Create parent stat directories in the single-task path too.

This branch still uses plain mkdir(), so nested targets like ./stat_files/model_1 work in multi-task mode and fail here.

Minimal fix

-                    Path(stat_file_path).mkdir()
+                    Path(stat_file_path).mkdir(parents=True, exist_ok=True)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/pt_expt/entrypoints/main.py` around lines 87 - 95, The code fails for
nested stat paths because it calls Path(...).mkdir() without parents=True;
update the branch that handles stat_file_path: when stat_file_path endswith
h5/hdf5, ensure its parent directory exists
(Path(stat_file_path).parent.mkdir(parents=True, exist_ok=True)) before creating
the h5 file, and when treating stat_file_path as a directory call
Path(stat_file_path).mkdir(parents=True, exist_ok=True); keep references to
stat_file_path and DPPath in the same block so the rest of the logic is
unchanged.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

deepmd/pt_expt/train/training.py (4)

812-818: ⚠️ Potential issue | 🟠 Major

Keep the same data_stat_protect when reapplying sharing after load.

The initial share_params() call uses the validated branch value, but the resume path drops it. Any non-default configuration will restore one state and then immediately re-share stats with different protection semantics.

Minimal fix

             if shared_links is not None:
                 # Re-apply sharing after loading checkpoint
                 self._unwrapped.share_params(
                     shared_links,
                     resume=True,
-                    model_key_prob_map=dict(zip(self.model_keys, self.model_prob)),
+                    model_key_prob_map=dict(
+                        zip(self.model_keys, self.model_prob, strict=True)
+                    ),
+                    data_stat_protect=_data_stat_protect[0],
                 )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/pt_expt/train/training.py` around lines 812 - 818, When reapplying
sharing in the resume path, include the same validated data_stat_protect value
so protection semantics aren’t lost; update the call to
self._unwrapped.share_params(...) inside the resume branch to pass the
data_stat_protect argument (the same variable used in the original share_params
call) along with shared_links, resume=True, and the model_key_prob_map
constructed from self.model_keys and self.model_prob.

---

499-508: ⚠️ Potential issue | 🟠 Major

Handle validation_data=None in the multi-task setup.

__init__() accepts validation_data: ... | None, but Line 505 indexes it unconditionally. A multi-task caller that disables validation will fail before training starts.

Minimal fix

             for model_key in self.model_keys:
                 data_requirement = self.loss[model_key].label_requirement
                 data_requirement += get_additional_data_requirement(
                     self.model[model_key]
                 )
                 training_data[model_key].add_data_requirements(data_requirement)
-                if validation_data[model_key] is not None:
-                    validation_data[model_key].add_data_requirements(data_requirement)
+                branch_validation = (
+                    None if validation_data is None else validation_data.get(model_key)
+                )
+                if branch_validation is not None:
+                    branch_validation.add_data_requirements(data_requirement)
                 self.training_data[model_key] = training_data[model_key]
-                self.validation_data[model_key] = validation_data[model_key]
+                self.validation_data[model_key] = branch_validation

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/pt_expt/train/training.py` around lines 499 - 508, The code
unconditionally indexes validation_data[model_key] which fails if the caller
passed validation_data=None; change the loop to guard that access—either replace
uses of validation_data[model_key] with a safe lookup (e.g., if validation_data
is not None and validation_data.get(model_key) is not None then call
add_data_requirements and assign self.validation_data[model_key] =
validation_data[model_key], otherwise set self.validation_data[model_key] =
None), or pre-normalize validation_data to a dict of Nones for all
self.model_keys before the loop so the existing logic (training_data,
add_data_requirements, and assigning self.validation_data) works without raising
when validation_data was None.

---

1139-1152: ⚠️ Potential issue | 🟠 Major

Don't consume training batches just to populate log output.

Line 1143 pulls extra training batches for the inactive heads on rank 0. That changes which samples those tasks actually train on, and in distributed runs only rank 0 advances those iterators further.

Please log from validation data or from a cached/non-advancing sample instead.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/pt_expt/train/training.py` around lines 1139 - 1152, The loop over
self.model_keys currently calls self.get_data(is_train=True, task_key=_key) and
self._unwrapped for inactive heads, which consumes training iterators and
advances only rank 0; instead, log using non-advancing data: replace those
get_data calls with a non-advancing source (e.g., use validation data via
get_data(is_train=False, task_key=_key) or a cached/sample snapshot you capture
earlier), and call self._unwrapped with that non-advancing label/input; keep the
rest of the collection logic (train_results[_key] = ...) the same so logging no
longer mutates training iterators for self.model_keys and task_key.

---

606-610: ⚠️ Potential issue | 🟡 Minor

Add strict=True to both zip() calls.

Both sites silently truncate if self.model_keys and self.model_prob ever drift, and Ruff will flag them either way.

Minimal fix

-                model_key_prob_map=dict(zip(self.model_keys, self.model_prob)),
+                model_key_prob_map=dict(
+                    zip(self.model_keys, self.model_prob, strict=True)
+                ),

-                    model_key_prob_map=dict(zip(self.model_keys, self.model_prob)),
+                    model_key_prob_map=dict(
+                        zip(self.model_keys, self.model_prob, strict=True)
+                    ),

As per coding guidelines, `**/*.py`: Install linter and run `ruff check .` before committing changes or the CI will fail.

Also applies to: 814-817

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@deepmd/pt_expt/train/training.py` around lines 606 - 610, The zip calls that
build the model_key_prob_map (pairing self.model_keys and self.model_prob) must
be made strict to avoid silent truncation; update the dict(zip(...)) in the
self.wrapper.share_params call (model_key_prob_map argument) to use
zip(self.model_keys, self.model_prob, strict=True), and do the same for the
other zip later in the file that also pairs self.model_keys with self.model_prob
so both mappings enforce equal lengths.

source/tests/pt_expt/test_multitask.py (2)

1951-1953: ⚠️ Potential issue | 🟡 Minor

Rename the unused n2 loop variable.

Ruff will flag n2 as unused here.

Minimal fix

-        for (n1, p1), (n2, p2) in zip(
+        for (n1, p1), (_, p2) in zip(
             mt_desc.named_parameters(), mt_desc_2.named_parameters(), strict=True
         ):

As per coding guidelines, `**/*.py`: Install linter and run `ruff check .` before committing changes or the CI will fail.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@source/tests/pt_expt/test_multitask.py` around lines 1951 - 1953, The loop
declares unused variable n2 causing linter errors; rename n2 to a throwaway name
(e.g., _n2 or _) in the for loop that iterates over mt_desc.named_parameters()
and mt_desc_2.named_parameters() so only p2 is treated as used (i.e., for (n1,
p1), (_n2, p2) in zip(..., strict=True)). Ensure the change touches the loop in
test_multitask.py where mt_desc and mt_desc_2 named_parameters() are zipped so
Ruff no longer flags an unused variable.

---

939-941: ⚠️ Potential issue | 🟡 Minor

Drop the unused finetune_links_true binding.

Ruff will flag this local as unused.

Minimal fix

-        model_config_true, finetune_links_true = get_finetune_rules(
+        model_config_true, _ = get_finetune_rules(
             ckpt_path, deepcopy(ft_config_true["model"]), change_model_params=True
         )

As per coding guidelines, `**/*.py`: Install linter and run `ruff check .` before committing changes or the CI will fail.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@source/tests/pt_expt/test_multitask.py` around lines 939 - 941, The call to
get_finetune_rules returns two values but the second (finetune_links_true) is
unused; remove the unused binding by either unpacking only the first return
(e.g., assign the single result of get_finetune_rules to model_config_true) or
use a throwaway name like _ for the second value so Ruff no longer flags
finetune_links_true as unused; update the invocation of get_finetune_rules where
model_config_true, finetune_links_true is currently used.

source/tests/pt_expt/test_training_ddp.py (3)

63-67: ⚠️ Potential issue | 🟠 Major

The free-port probe is still racy for DDP rendezvous.

This helper releases the socket before the workers call init_process_group(), so another parallel test can claim the same port and make the suite flaky. Please switch to a rendezvous that stays allocated for the whole run, e.g. file://... or a pre-created TCPStore.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@source/tests/pt_expt/test_training_ddp.py` around lines 63 - 67, The helper
_find_free_port() is racy because it closes the probe socket before workers call
init_process_group(), allowing races on the chosen port; replace this pattern by
using a rendezvous that stays allocated for the whole test run (for example
switch the rendezvous URL from tcp://<port> to a file://... path or create a
persistent torch.distributed.TCPStore ahead of spawning workers) and update the
test harness that builds the init_process_group() args to use the file-based
rendezvous or the pre-created TCPStore so the address/resource remains reserved
for the duration of the DDP test.

---

17-18: ⚠️ Potential issue | 🟠 Major

Set DEVICE=cpu before importing deepmd.pt_expt.

The worker-level override is too late here. mp.spawn starts a fresh interpreter that imports this module before entering _worker_*, so the module-scope deepmd.pt_expt imports can lock in the parent/default device first. On CUDA-capable runners that makes this “CPU-only” suite depend on the ambient device instead of the explicit test setting.

Suggested fix

 import os
 import shutil
 import socket
 import tempfile
 import unittest
@@
+os.environ["DEVICE"] = "cpu"
+
 from deepmd.pt_expt.entrypoints.main import (
     get_trainer,
 )

Also applies to: 34-48

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@source/tests/pt_expt/test_training_ddp.py` around lines 17 - 18, Set the
environment variable DEVICE=cpu before importing deepmd.pt_expt in the module
scope of the test file so spawned worker interpreters don't inherit the ambient
CUDA device; modify the top of source/tests/pt_expt/test_training_ddp.py to
export DEVICE=cpu (e.g., via os.environ["DEVICE"]="cpu") prior to any import
that pulls in deepmd.pt_expt or before any use of mp.spawn/_worker_* so the
child processes import with the correct CPU-only setting.

---

623-632: ⚠️ Potential issue | 🟠 Major

These spawned training tests still need a hard 60s timeout.

mp.spawn(..., join=True) can hang indefinitely on a broken rendezvous or stuck worker, and this pattern repeats through the rest of the module. Please wrap the spawn call in a bounded subprocess/helper so each training case fails fast after 60 seconds.

As per coding guidelines "Set training test timeouts to 60 seconds maximum for validation purposes, as real training takes hours or days".

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@source/tests/pt_expt/test_training_ddp.py` around lines 623 - 632, The test
test_ddp_single_task_trains uses mp.spawn to run _worker_single_task_train and
can hang indefinitely; wrap the mp.spawn call in a bounded timeout helper (e.g.,
run it in a short-lived subprocess or use a helper that starts the spawn in a
separate Process and joins with a 60-second timeout) so that if the spawn does
not complete within 60 seconds the test forcefully terminates the worker(s) and
fails; ensure the same pattern (60s hard timeout) is applied wherever mp.spawn
is called in this module.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes